#ifndef __SPLINE__
#define __SPLINE__

#include <iostream>
#include <vector>
#include <set>
#include "Function.h"
#include "Hermite.h"

#include <Eigen/Sparse>
#include <Eigen/SparseLU>

using namespace std;
using namespace Eigen;

enum BCType {
    Complete,
    Second,
    Natural,
    None //For spline without boundary condition
};

class Spline {
public:
    virtual void solve() = 0;
    virtual double operator()(const double &point) = 0;
};

//N-order piecewise polynomial Spline
template <int N, BCType bctype = None>
class ppformSpline : public Spline {
private:
    int n; 
    const vector<double>&x, &y;
    BCType bct;
    vector<Polynomial> polys;
    const double& cond_left;
    const double& cond_right;

public:
    ppformSpline(const vector<double>& _x, const vector<double>& _y,
                 const double& _cond_left = 0, const double& _cond_right = 0)
        : n(_x.size()), x(_x), y(_y), cond_left(_cond_left), cond_right(_cond_right), bct(bctype) {
        if (_x.size() != _y.size()) {
            cerr << "Error: (ppspline) the size of x and y should be the same." << endl;
            exit(-1);
        }
        solve();
    }

    
    void solve() {
        if (N == 1) {
            for (int i = 1; i < n; i++) {
                // p(x) = k * x + b
                double k = (y[i] - y[i - 1]) / (x[i] - x[i - 1]);
                double b = -x[i - 1] * k + y[i - 1];
                polys.push_back(Polynomial(b, k));
            }
        } else if (N == 3) {
            if (bct == Complete) {
                SparseMatrix<double> A(n - 2, n - 2);
                VectorXd b(n - 2), r(n - 2);
                vector<Triplet<double>> tripletlist;

                for (int i = 0; i < n - 2; i++) {
                    tripletlist.push_back(Triplet<double>(i, i, 2));
                    if (i != 0) {
                        double lambda = (x[i + 2] - x[i + 1]) / (x[i + 2] - x[i]);
                        tripletlist.push_back(Triplet<double>(i, i - 1, lambda));
                    }
                    if (i != n-3) {
                        double mu = (x[i + 1] - x[i]) / (x[i + 2] - x[i]);
                        tripletlist.push_back(Triplet<double>(i, i + 1, mu));
                    }
                    b(i) = 3 * (x[i + 1] - x[i]) / (x[i + 2] - x[i]) * (y[i + 2] - y[i + 1]) / (x[i + 2] - x[i + 1]) + 3 * (x[i + 2] - x[i + 1]) / (x[i + 2] - x[i]) * (y[i + 1] - y[i]) / (x[i + 1] - x[i]);
                }
                A.setFromTriplets(tripletlist.begin(), tripletlist.end());
                A.makeCompressed();
                b(0) -= (x[2] - x[1]) / (x[2] - x[0]) * cond_left;
                b(n - 3) -= (x[n - 2] - x[n - 3]) / (x[n - 1] - x[n - 3]) * cond_right;
                
                SparseLU<SparseMatrix<double>> solver;
                solver.compute(A);
                r = solver.solve(b);
        
                VectorXd m(n);
                m(0) = cond_left;
                m(n - 1) = cond_right;
                m.segment(1, n - 2) = r;
                for (int i = 0; i < n - 1; i++) {
                    vector<double> x_h = { x[i], x[i], x[i + 1], x[i + 1] };
                    vector<vector<double>> y_h = { { y[i], y[i], y[i + 1], y[i + 1] }, { m[i], m[i], m[i + 1], m[i + 1] } };
                    Hermite herm(x_h, y_h);
                    polys.push_back(herm.get_poly());
                }

            } else if (bct == Second || bct == Natural) {
                SparseMatrix<double> A(n, n);
                VectorXd b(n), r(n);
                vector<Triplet<double>> tripletlist;
                for (int i = 0; i < n; i++) {
                    tripletlist.push_back(Triplet<double>(i, i, 2));
                    if (i != 0 && i != n - 1) {
                        double lambda = (x[i + 1] - x[i]) / (x[i + 1] - x[i - 1]);
                        double mu = (x[i] - x[i - 1]) / (x[i + 1] - x[i - 1]);
                        tripletlist.push_back(Triplet<double>(i, i - 1, lambda));
                        tripletlist.push_back(Triplet<double>(i, i + 1, mu));
                        b(i) = 3 * (x[i] - x[i - 1]) / (x[i + 1] - x[i - 1]) * (y[i + 1] - y[i]) / (x[i + 1] - x[i]) + 3 * (x[i + 1] - x[i]) / (x[i + 1] - x[i - 1]) * (y[i] - y[i - 1]) / (x[i] - x[i - 1]);
                    } 
                }
                tripletlist.push_back(Triplet<double>(n - 1, n - 2, 1));
                tripletlist.push_back(Triplet<double>(0, 1, 1));
                A.setFromTriplets(tripletlist.begin(), tripletlist.end());
                A.makeCompressed();
                b(0) = 3 * (y[1] - y[0]) / (x[1] - x[0]) - (x[1] - x[0]) / 2 * cond_left;
                b(n - 1) = 3 * (y[n - 1] - y[n - 2]) / (x[n - 1] - x[n - 2]) + (x[n - 1] - x[n - 2]) / 2 * cond_right;

                SparseLU<SparseMatrix<double>> solver;
                solver.compute(A);
                r = solver.solve(b);

                VectorXd m = r;
                for (int i = 0; i < n - 1; i++) {
                    vector<double> x_h = { x[i], x[i], x[i + 1], x[i + 1] };
                    vector<vector<double>> y_h = { { y[i], y[i], y[i + 1], y[i + 1] }, { m[i], m[i], m[i + 1], m[i + 1] } };
                    Hermite herm(x_h, y_h);
                    polys.push_back(herm.get_poly());
                }
            } else {
                cerr << "Error: cannot compute the boundary condition." << endl;
                exit(-1);
            }
        } else {
            cerr << "Error: cannot compute the order." << endl;
            exit(-1);
        }
    };

    
    double operator()(const double &point) {
        int pos = 0;
        while (point > x[pos + 1]) {
            pos++;
        }
        return polys[pos](point);
    }
};




//N-order B-spline
template <int N, BCType bctype = None>
class Bspline : public Spline {
private:
    int n;
    const vector<double>&x, y;
    vector<double> x_complete;
    BCType bct;
    vector<BsplineBase> bases;
    VectorXd coef;
    const double& cond_left;
    const double& cond_right;

public:
    Bspline(const vector<double>& _x, const vector<double>& _y,
            const double& _cond_left = 0, const double& _cond_right = 0)
        : n(_x.size()), y(_y), x(_x), x_complete(_x), cond_left(_cond_left), cond_right(_cond_right), bct(bctype) {
        
        //fill x_complete
        vector<double> x_left, x_right;
        double d = x[1] - x[0];
        if (N == 1) {
            x_left = { x[0] - d };
            x_right = { x[n - 1] + d };
        } else if (N == 2) {
            x_left = { x[0] - 2 * d, x[0] - d };
            x_right = { x[n - 1] + d, x[n - 1] + 2 * d };
        } else if (N == 3) {
            x_left = { x[0] - 3 * d, x[0] - 2 * d, x[0] - d };
            x_right = { x[n - 1] + d, x[n - 1] + 2 * d, x[n - 1] + 3 * d };
        } else {
            cerr << "Error: cannot compute the order." << endl;
            exit(-1);
        }
        x_complete.insert(x_complete.begin(), x_left.begin(), x_left.end());
        x_complete.insert(x_complete.end(), x_right.begin(), x_right.end());

        //fill bases
        for (int i = 2 - N; i < n + 1; i++) {
            vector<double> knots(x_complete.begin() + i - 2 + N, x_complete.begin() + i + 2*N);
            BsplineBase base(N, knots);
            bases.push_back(base);
        }
    }

    
    void solve() {
        if (N == 1) {
            coef.resize(n);
            for (int i = 0; i < n; i++) {
                coef[i] = y[i];
            }
        } else if (N == 2) {
            SparseMatrix<double> A(n - 1, n - 1);
            VectorXd b(n - 1), r(n - 1);
            vector<Triplet<double>> tripletlist;

            for (int i = 0; i < n - 1; i++) {
                if (i != 0) {
                    tripletlist.push_back(Triplet<double>(i, i - 1, 1));
                    tripletlist.push_back(Triplet<double>(i - 1, i, 1));
                }
                if (i != 0 && i != n - 2) {
                    tripletlist.push_back(Triplet<double>(i, i, 6));
                }
                b[i] = 8 * y[i];
            }
            tripletlist.push_back(Triplet<double>(0, 0, 5));
            tripletlist.push_back(Triplet<double>(n - 2, n - 2, 5));
            A.setFromTriplets(tripletlist.begin(), tripletlist.end());
            A.makeCompressed();
            b[0] = b[0] - 2 * cond_left;
            b[n - 2] = b[n - 2] - 2 * cond_right;

            SparseLU<SparseMatrix<double>> solver;
            solver.compute(A);
            coef.resize(n + 1);
            r = solver.solve(b);
            coef[0] = 2 * cond_left - r[0];
            coef[n] = 2 * cond_right - r[n - 2];
            coef.segment(1, n - 1) = r;

        } else if (N == 3) {
            if (bct == Complete) {
                SparseMatrix<double> A(n + 2, n + 2);
                VectorXd b(n + 2), r(n + 2);
                vector<Triplet<double>> tripletlist;

                for (int i = 1; i < n + 1; i++) {
                    for (int j = i - 1; j < i + 2; j++) {
                        double tmp = bases[j](x_complete[i + 2]);
                        tripletlist.push_back(Triplet<double>(i, j, tmp));
                    }
                }

                vector<BsplineBase> bases2;
                for (int i = 0; i < n + 1; i++) {
                    vector<double> knots2(x_complete.begin() + i + 1, x_complete.begin() + i + 5);
                    BsplineBase base2(2, knots2);
                    bases2.push_back(base2);
                }
                double tmp1 = -3 * bases2[0](x_complete[3]) / (x_complete[4] - x_complete[1]);
                double tmp2 = 3 * bases2[1](x_complete[3]) / (x_complete[5] - x_complete[2]);
                double tmp3 = -3 * bases2[n - 1](x_complete[n + 2]) / (x_complete[n + 3] - x_complete[n]);
                double tmp4 = 3 * bases2[n](x_complete[n + 2]) / (x_complete[n + 4] - x_complete[n + 1]);
                tripletlist.push_back(Triplet<double>(0, 0, tmp1));
                tripletlist.push_back(Triplet<double>(0, 1, - tmp1 - tmp2));
                tripletlist.push_back(Triplet<double>(0, 2, tmp2));
                tripletlist.push_back(Triplet<double>(n + 1, n - 1, tmp3));
                tripletlist.push_back(Triplet<double>(n + 1, n, - tmp3 - tmp4));
                tripletlist.push_back(Triplet<double>(n + 1, n + 1, tmp4));

                A.setFromTriplets(tripletlist.begin(), tripletlist.end());
                A.makeCompressed();
                b[0] = cond_left;
                b[n + 1] = cond_right;
                for (int i = 1; i < n + 1; i++) {
                    b[i] = y[i - 1];
                }

                SparseLU<SparseMatrix<double>> solver;
                solver.compute(A);
                coef = solver.solve(b);

            } else if (bct == Second || bct == Natural) {
                SparseMatrix<double> A(n + 2, n + 2);
                VectorXd b(n + 2), r(n + 2);
                vector<Triplet<double>> tripletlist;

                for (int i = 1; i < n + 1; i++) {
                    for (int j = i - 1; j < i + 2; j++) {
                        double tmp = bases[j](x_complete[i + 2]);
                        tripletlist.push_back(Triplet<double>(i, j, tmp));
                    }
                }
                double tmp1 = 6.0 / ((x_complete[4] - x_complete[1]) * (x_complete[4] - x_complete[2]));
                double tmp2 = 6.0 / ((x_complete[5] - x_complete[2]) * (x_complete[4] - x_complete[2]));
                double tmp3 = 6.0 / ((x_complete[n + 3] - x_complete[n]) * (x_complete[n + 3] - x_complete[n + 1]));
                double tmp4 = 6.0 / ((x_complete[n + 4] - x_complete[n + 1]) * (x_complete[n + 3] - x_complete[n + 1]));
                tripletlist.push_back(Triplet<double>(0, 0, tmp1));
                tripletlist.push_back(Triplet<double>(0, 1, -tmp1 - tmp2));
                tripletlist.push_back(Triplet<double>(0, 2, tmp2));
                tripletlist.push_back(Triplet<double>(n + 1, n - 1, tmp3));
                tripletlist.push_back(Triplet<double>(n + 1, n, -tmp3 - tmp4));
                tripletlist.push_back(Triplet<double>(n + 1, n + 1, tmp4));
        
                A.setFromTriplets(tripletlist.begin(), tripletlist.end());
                A.makeCompressed();
                b[0] = cond_left;
                b[n + 1] = cond_right;
                for (int i = 1; i < n + 1; i++) {
                    b[i] = y[i - 1];
                }

                SparseLU<SparseMatrix<double>> solver;
                solver.compute(A);
                coef = solver.solve(b);
            }
        }
    }

    double operator()(const double &point) {
        solve();
        if (N == 1) {
            int pos = 1;
            while (point > x_complete[pos + 1]) {
                pos++;
            }
            return coef[pos - 1] * bases[pos - 1](point) + coef[pos] * bases[pos](point);
        } else if (N == 2) {
            int pos = 1;
            while (point > x_complete[pos + 2]) {
                pos++;
            }
            return coef[pos - 1] * bases[pos - 1](point) + coef[pos] * bases[pos](point) + coef[pos + 1] * bases[pos + 1](point);
        } else if (N == 3) {
            int pos = 1;
            while (point > x_complete[pos + 3]) {
                pos++;
            }
            return coef[pos - 1] * bases[pos - 1](point) + coef[pos] * bases[pos](point) + coef[pos + 1] * bases[pos + 1](point) + coef[pos + 2] * bases[pos + 2](point);
        }
    }
};

#endif
